# continue_hpo.py
import os
import sys
import glob
import sqlite3
import argparse
import json

# 添加项目根目录到Python路径
ROOT_DIR = '/media/ross/8TB/project/lsh/deep_learning/DiffusionDet_mmdet/DiffusionDet'
sys.path.insert(0, ROOT_DIR)

from tools.mymodel_analysis_tools.HPO.train_hpo import HPOptimization


def get_latest_study_info():
    """获取最新的研究信息"""
    # 查找所有的HPO数据库文件
    db_files = glob.glob("hpo_*.db")

    if not db_files:
        print("未找到任何HPO数据库文件，请先运行一次HPO优化。")
        return None

    # 按修改时间排序
    db_files.sort(key=lambda x: os.path.getmtime(x), reverse=True)
    latest_db = db_files[0]

    # 从文件名中提取时间戳
    timestamp = latest_db.replace("hpo_", "").replace(".db", "")
    study_name = f"diffusiondet_hpo_{timestamp}"

    # 连接数据库并获取已完成的试验数
    conn = sqlite3.connect(latest_db)
    cursor = conn.cursor()
    cursor.execute("SELECT COUNT(*) FROM trials WHERE state = 'COMPLETE'")
    completed_trials = cursor.fetchone()[0]
    conn.close()

    return timestamp, study_name, latest_db, completed_trials


def check_trial_status(db_path):
    """检查试验状态"""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT state, COUNT(*) FROM trials GROUP BY state")
    states = cursor.fetchall()
    conn.close()

    print("试验状态统计:")
    for state, count in states:
        print(f"  {state}: {count}")


def main():
    parser = argparse.ArgumentParser(description='继续HPO优化')
    parser.add_argument('--config',
                        default="work_dirs/Hyperparametric_sensitivity/diffusiondet_r50_fpn_epoch_microalgeaOri20%_1lcm2_1adem2_1ddim4_1distill4.py",
                        help='配置文件路径')
    parser.add_argument('--work-dir', default="work_dirs/Hyperparametric_sensitivity/hpo_experiment",
                        help='工作目录')
    parser.add_argument('--n-trials', type=int, default=25,
                        help='总共要运行的试验数')
    parser.add_argument('--timestamp', default=None,
                        help='指定时间戳，如果不指定则使用最新的')
    parser.add_argument('--weights', default=None,
                        help='目标函数权重JSON文件路径')
    args = parser.parse_args()

    # 配置基本参数
    base_config = os.path.join(ROOT_DIR, args.config)
    work_dir = os.path.join(ROOT_DIR, args.work_dir)

    # 加载目标函数权重
    objective_weights = None
    if args.weights:
        with open(args.weights, 'r') as f:
            objective_weights = json.load(f)
        print(f"使用自定义目标函数权重: {objective_weights}")

    # 获取最新的研究信息
    if args.timestamp:
        timestamp = args.timestamp
        study_name = f"diffusiondet_hpo_{timestamp}"
        storage = f"sqlite:///hpo_{timestamp}.db"
        print(f"使用指定的时间戳: {timestamp}")
        check_trial_status(f"hpo_{timestamp}.db")
    else:
        result = get_latest_study_info()
        if result is None:
            print("未找到任何HPO数据库文件，请先运行一次HPO优化或指定时间戳。")
            return

        timestamp, study_name, storage_file, completed_trials = result
        storage = f"sqlite:///{storage_file}"
        print(f"找到最新的研究: {study_name}")
        print(f"已完成的试验数: {completed_trials}")
        print(f"数据库文件: {storage_file}")
        check_trial_status(storage_file)

    print(f"使用配置文件: {base_config}")
    print(f"工作目录: {work_dir}")

    # 创建优化器
    optimizer = HPOptimization(
        base_config_path=base_config,
        work_dir=work_dir,
        n_trials=args.n_trials,
        study_name=study_name,
        storage=storage,
        objective_weights=objective_weights,
        resume=True
    )

    # 运行优化
    optimizer.run_optimization()


if __name__ == "__main__":
    main()
